package cn.njupt.tree;

import cn.njupt.util.Comparator;

//AVL树
public class AVLTree<E> extends BST<E> {


    public AVLTree() {
        this(null);
    }

    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    @Override
    protected void afterAdd(Node<E> node) {
        while((node=node.parent)!=null){
            if(isBalance(node)){ //看node是否平衡
                //更新高度
                updateHeight(node);
            }else{
                //恢复平衡
                reBalance2(node);   //node代表的是grand
                break;
            }
        }
    }

    @Override
    protected void afterRemove(Node<E> node) {
        while((node=node.parent)!=null){
            if(isBalance(node)){
                updateHeight(node);
            }else{
                reBalance(node);
            }
        }
    }

    /**
     * 恢复平衡
     * @param grand
     */
    private void reBalance(Node<E> grand){
        Node<E> parent=((AVLNode<E>)grand).tallerChild();
        Node<E> child=((AVLNode<E>)parent).tallerChild();
        if(parent.isLeftChild()){
            if(child.isLeftChild()){
                //LL   做法：右旋转
                rotateRight(grand);
            }else{
                //LR   做法：parent先左旋转，grand再右旋转
                rotateLeft(parent);
                rotateRight(grand);
            }
        }else{
            if(child.isLeftChild()){
                //RL   做法：parent先右旋转，grand再左旋转
                rotateRight(parent);
                rotateLeft(grand);
            }else{
                //RR   做法：左旋转
                rotateLeft(grand);
            }
        }
    }

    /**
     * 恢复平衡，利用统一旋转
     * @param grand
     */
    private void reBalance2(Node<E> grand){
        Node<E> parent=((AVLNode<E>)grand).tallerChild();
        Node<E> child=((AVLNode<E>)parent).tallerChild();
        if(parent.isLeftChild()){
            if(child.isLeftChild()){
                //LL   做法：右旋转
                rotate(grand,child.left,child,child.right,parent,parent.right,grand,grand.right);
            }else{
                //LR   做法：parent先左旋转，grand再右旋转
               rotate(grand,parent.left,parent,child.left,child,child.right,grand,grand.right);
            }
        }else{
            if(child.isLeftChild()){
                //RL   做法：parent先右旋转，grand再左旋转
                rotate(grand,grand.left,grand,child.left,child,child.right,parent,parent.right);
            }else{
                //RR   做法：左旋转
                rotate(grand,grand.left,grand,parent.left,parent,child.left,child,child.right);
            }
        }
    }

    /**
     * 利用统一的规则进行统一旋转
     */
    public void rotate(Node<E> root,//子树的根节点
                           Node<E> a,Node<E> b,Node<E> c,
                            Node<E> d,
                           Node<E> e,Node<E> f,Node<E> g){

        d.parent=root.parent;

        if(root.isLeftChild()){
            root.parent.left=d;
        }else if(root.isRightChild()){
            root.parent.right=d;
        }else{
            this.root=d;
        }
        b.left=a;
        if(a!=null){
            a.parent=b;
        }
        b.right=c;
        if(c!=null){
            c.parent=b;
        }

        updateHeight(b);
        f.left=e;
        if(e!=null){
            e.parent=f;
        }

        f.right=g;
        if(g!=null){
            g.parent=f;
        }
        updateHeight(f);

        b.parent=d;
        f.parent=d;
        d.left=b;
        d.right=f;

        updateHeight(d);

    }

    /**
     * 左旋转
     * @param grand
     */
    public void rotateLeft(Node<E> grand){
        //node节点肯定是平衡因子最大的那个节点
        Node<E> parent=grand.right;
        Node<E> child=parent.left;

        grand.right=child;
        parent.left=grand;

        parent.parent=grand.parent;
        if(grand.isLeftChild()){
            grand.parent.left=parent;
        }else if(grand.isRightChild()){
            grand.parent.right=parent;
        }else{
            root=parent;
        }
        if(child!=null){
            child.parent=grand;
        }
        grand.parent=parent;
        //更新高度
        updateHeight(grand);
        updateHeight(parent);
    }

    /**
     * 右旋转
     * @param grand
     */
    public void rotateRight(Node<E> grand){
        Node<E> parent=grand.left;
        Node<E> child=parent.right;

        grand.left=child;
        parent.right=grand;

        parent.parent=grand.parent;
        //判断parent的位置是父亲节点的左子树还是右子树
        if(grand.isLeftChild()){
            grand.parent.left=parent;
        }else if(grand.isRightChild()){
            grand.parent.right=parent;
        }else{
            root=parent;
        }

        if(child!=null){
            child.parent=grand;
        }
        grand.parent=parent;


        updateHeight(grand);
        updateHeight(parent);


    }


    /**
     * 更新高度
     * @param node
     */
    private void updateHeight(Node<E> node){
        ((AVLNode<E>)node).updateHeight();
    }

    private boolean isBalance(Node<E> node){
        AVLNode<E> avlNode=((AVLNode<E>)node);
        return Math.abs(avlNode.balanceFactor())<= 1;
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element,parent);
    }

    private static class AVLNode<E> extends Node<E> {
        int height = 1;

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        //计算当前节点的平衡因子
        public int balanceFactor(){
            int leftHeight=this.left==null?0:((AVLNode<E>)left).height;
            int rightHeight=this.right==null?0:((AVLNode<E>)right).height;
            return leftHeight-rightHeight;
        }

        //更新当前节点高度
        public void updateHeight(){
            int leftHeight=this.left==null?0:((AVLNode<E>)left).height;
            int rightHeight=this.right==null?0:((AVLNode<E>)right).height;
            this.height=Math.max(leftHeight,rightHeight)+1;
        }
        //返回当前节点更高的孩子节点
        public Node<E> tallerChild(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            if(leftHeight>rightHeight)return left;
            if(leftHeight<rightHeight)return right;
            return this.isLeftChild()?left:right;
        }

        @Override
        public String toString() {
            String str=null;
            if(parent!=null){
                str=parent.element.toString();
            }

            return this.element+"_p:"+str+"height="+this.height;
        }
    }
}
